#include <unordered_map>
#include <unistd.h>
#include <algorithm>
#include <math.h>
#include <numeric>
#include <thread>
#include <vector>
#include <memory>
#include <sys/time.h>
#include <random>
#include <faiss/ascend/custom/AscendIndexIVFSQT.h>

namespace {
    typedef std::unordered_map<int, float> recallMap;

template<class T>
recallMap calRecallNomal(std::vector<T> label, int64_t* gt, int shape)
{
    recallMap recMap;
    recMap[1] = 0;
    recMap[10] = 0;
    recMap[100] = 0;
    int k = label.size() / shape;
    for (int i = 0; i <  shape; i++) {
        for (int j = 0; j < k; j++) {
            if (gt[i * k] == label[i * k + j]) {
                recMap[100]++;
                switch (j)
                {
                    case 0:
                        recMap[1]++;
                        recMap[10]++;
                        break;
                    case 1 ... 9:
                        recMap[10]++;
                        break;
                    default:
                        break;
                }
                break;
            }
        }
    }
    recMap[1] = recMap[1] / shape * 100;
    recMap[10] = recMap[10] / shape * 100;
    recMap[100] = recMap[100] / shape * 100;
    return recMap;
}

inline double getMillisecs()
{
    struct timeval tv = {0, 0};
    gettimeofday(&tv, nullptr);
    return tv.tv_sec * 1e3 + tv.tv_usec * 1e-3;
}

int8_t randomInt8()
{
    int r = rand() % 256;
    uint8_t u = static_cast<uint8_t>(r);
    int8_t i = static_cast<int8_t>(u ^ 0x80);
    return i;
}

struct dataFloat
{
    dataFloat( std::vector<float> &base,
    std::vector<float> &learn,
    std::vector<float> &query,
    std::vector<int64_t> &gt)
    : base(base), learn(learn), query(query), gt(gt) {};
    std::vector<float> &base;
    std::vector<float> &learn;
    std::vector<float> &query;
    std::vector<int64_t> &gt;
};


void searchProccess(int dimIn, int nprobe, int l2probe, int l3probe, int topk, 
                    dataFloat data, faiss::ascend::AscendIndexIVFSQT &index)
{
    int flatK = 1;
    index.updateTParams(l2probe, l3probe);
    index.setNumProbes(nprobe);
    std::vector<size_t> searchNums{10240, 40960, 163840, 1310720};
    for (auto bs: searchNums) {
        double searchStart = getMillisecs();
        std::vector<float> dist(bs * static_cast<size_t>(topk), 0);
        std::vector<faiss::idx_t> labelTopk(bs * static_cast<size_t>(topk), 0);
        std::vector<faiss::idx_t> labelFlat(bs * static_cast<size_t>(topk), 0);
        std::vector<float> distanceTopk(bs * static_cast<size_t>(topk), 0);
        std::vector<faiss::idx_t> labelRet(bs * static_cast<size_t>(topk), 0);
        std::vector<float> distanceFlat(bs * static_cast<size_t>(flatK), 0);
        double sqtSearchStart = getMillisecs();
        index.search(bs, data.query.data(), topk, dist.data(), labelTopk.data());
        double sqtSearchEnd = getMillisecs();
        const int CPU_NUM = 48;
        auto cpuSearchFunctor = [&] (int cpuIdx, int start, int end) {
            for (size_t qId = start; qId < std::min(static_cast<size_t>(end), bs); qId++) {
                cpu_set_t cpuset;
                CPU_ZERO(&cpuset);
                CPU_SET(cpuIdx, &cpuset);
                float max = 0;
                size_t offset = 0;
                for (size_t topkId = 0; topkId <  static_cast<size_t>(topk); topkId++) {
                    bool isValid = false;
                    if (labelTopk[qId * topk + topkId] == -1) {
                        continue;
                    }
                    for (size_t dimId = 0; dimId < static_cast<size_t>(dimIn); dimId++) {
                        //计算IP距离
                        distanceTopk[qId * topk + topkId] +=
                        data.query[qId * dimIn + dimId] * \
                        data.base[labelTopk[qId * topk + topkId] * dimIn + dimId];
                        isValid = true;
                    }
                    if (isValid && max < distanceTopk[qId * topk + topkId]) {
                        max = distanceTopk[qId * topk + topkId];
                        offset = topkId;
                    }
                    labelRet[qId * topk] = labelTopk[qId * topk + offset];
                    distanceFlat[qId * flatK] = distanceTopk[qId * topk + offset];
                }
            }
        };
        std::thread threads[CPU_NUM];
        int block = (bs + CPU_NUM - 1) / CPU_NUM; // div up
        for (int i = 0; i < CPU_NUM; i++) {
            threads[i] = std::thread(cpuSearchFunctor, i, i * block, (i + 1) * block);
        }
        for (int i = 0; i < CPU_NUM; i++) {
            threads[i].join();
        }
        double cpuSearchEnd = getMillisecs();
        recallMap recall = calRecallNomal(labelRet, data.gt.data(), bs);
        printf("bs = %zu, r1 = %.2f, r10 = %.2f, r100 = %.2f, sqt search cost: %lf, cpu search cost: %lf, qps = %lf\n",
                bs, recall[1], recall[10], recall[100], sqtSearchEnd - sqtSearchStart,
                cpuSearchEnd - sqtSearchStart, 1000 * bs / (cpuSearchEnd - searchStart));

    }
}

void addAndSearch(int niter, int fuzzy_k, int ncentroids, float threshold, std::vector<int> nprobes,
                  std::vector<int> l2probes, std::vector<int> l3probes, int trainSize, int topk)
{
    int dimIn = 256;
    int dimOut = 64;
    size_t ntotal = 10000000;
    size_t queryNum = 10000000;
    size_t learnNum = ntotal / 10;
    size_t dim = 256;

    std::vector<int> devices = {0, 1, 2, 3};
    int64_t resourceSize = 1 * static_cast<int64_t>(1024 * 1024 * 1024);
    faiss::ascend::AscendIndexIVFSQTConfig conf(devices, resourceSize);
    conf.cp.niter = niter;
    conf.useKmeansPP = true;
    conf.cp.max_points_per_centroid = 256;
    faiss::ascend::AscendIndexIVFSQT index(dimIn, dimOut, ncentroids, faiss::ScalarQuantizer::QuantizerType::QT_8bit,
                                           faiss::METRIC_INNER_PRODUCT, conf);
    printf("start generate data\n");
    std::vector<int8_t> baseInt8(ntotal * dim);
    for (size_t i = 0; i < ntotal * dim; i++) {
        baseInt8[i] = randomInt8();
    }
    std::vector<int8_t> learnInt8(learnNum * dim);
    for (size_t i = 0; i < learnNum * dim; i++) {
        learnInt8[i] = randomInt8();
    }
    std::vector<int8_t> queryInt8(queryNum * dim);
    std::vector<int64_t> gt(queryNum * 100, 0);
    for (size_t q = 0; q < queryNum; q++) {
        size_t idx = static_cast<size_t>(rand() % ntotal);
        for (size_t d = 0; d < dim; d++) {
            queryInt8[q * dim + d] = baseInt8[idx * 256 + d];
        }
        gt[q * 100] = static_cast<int64_t>(idx);
    }


    //int8 to float
    std::vector<float> base(ntotal * dim);
    for (size_t i = 0; i < ntotal * dim; i++) {
        base[i] = static_cast<float>(baseInt8[i]) / 128.0;
    }
    std::vector<float> query(queryNum * dim);
    for (size_t i = 0; i < queryNum * dim; i++) {
        query[i] = static_cast<float>(queryInt8[i]) / 128.0;
    }
    std::vector<float> learn(learnNum * dim);
    for (size_t i = 0; i < learnNum * dim; i++) {
        learn[i] = static_cast<float>(learnInt8[i]) / 128.0;
    }
    dataFloat dataFloat(base, learn, query, gt);
    printf("generate data ok\n");
    printf("topk num is:%d\n", topk);
    index.verbose = true;
    index.setFuzzyK(fuzzy_k);
    index.setThreshold(threshold);
    if (learnNum < 6000000) {
        trainSize = learnNum;
    }
    double trainStart = getMillisecs();
    index.train(trainSize, learn.data());
    double trainEnd = getMillisecs();
    printf("train time cost : %.2fs\n", (trainEnd - trainStart) / 1000);
    double addStart = getMillisecs();
    index.add(ntotal, base.data());
    double addEnd = getMillisecs();
    printf("add time cost : %.2fs\n", (addEnd - addStart) / 1000);
    double updateStart = getMillisecs();
    index.update();
    double updateEnd = getMillisecs();
    printf("update time cost : %.2fs\n", (updateEnd - updateStart) / 1000);
    printf("=> start Qps Test!\n");
    for (int nprobe: nprobes) {
        for (int l2: l2probes) {
            for (int l3: l3probes) {
                searchProccess(dimIn, nprobe, l2, l3, topk, dataFloat, index);
            }
        }
    }

}

} // namespace

int main(int argc, char **argv)
{
    float threshold = 1.6; // 1.6
    int ncentroids = 16384; // 分桶的数目
    int topk = 100;
    int trainSize = 5000000;
    printf("Start Test\n");
    std::vector<int> nprobes = {64};
    std::vector<int> l2probes = {360};
    std::vector<int> l3probes = {360};
    addAndSearch(16, 3, ncentroids, threshold, nprobes, l2probes, l3probes, trainSize, topk);
}